from random import seed
import numpy as np
import torch
from pytorch_lightning.callbacks import ModelCheckpoint
from model import MyModel
from dataset import MyData
import config
import pytorch_lightning as pl

opt = config.get_options()

def set_seed(seed_value):
    seed(seed_value)
    np.random.seed(seed_value)
    torch.manual_seed(seed_value)
    torch.cuda.manual_seed_all(seed_value)

checkpoint_callback = ModelCheckpoint(
    monitor='val_loss',
    filename='sample-mnist-{epoch:02d}-{val_loss:.2f}',
    save_top_k=3,
    mode='min',
    save_last=True
)

if __name__ == '__main__':
    set_seed(opt.seed)
    model = MyModel(
        input_size=opt.input_size,
        learning_rate=opt.lr,
        hidden_size=opt.hidden_size,
        dropout_rate=opt.dropout_rate
    )

    dm = MyData(
        data_dir = opt.data_dir, 
        batch_size=opt.batch_size, 
        num_workers=opt.workers
    )

    trainer = pl.Trainer(
        min_epochs=1,
        max_epochs=opt.niter, 
        callbacks=[checkpoint_callback],
        logger=True
    )
    trainer.fit(model, dm)
    trainer.validate(model, dm)
    trainer.test(model, dm)